from google.cloud import documentai_v1 as documentai
from google.protobuf.json_format import MessageToDict
import json
import cv2
import numpy as np
import os
from utils import fsutils

# Replace with your values
project_id = "fourth-amp-399005"
location = "us"  # e.g., "us" or "eu"
processor_id = "305ba17286ecb089"


# 获取原始 OCR 结果的文件路径
def get_ocr_raw_result_path(filepath):
    # 获取用于存储原始 OCR 结果的父目录路径
    parent_dir = fsutils.get_cache_path_for_category("ocr_raw_result")
    # 构建最终的原始 OCR 结果文件路径
    rrpath = os.path.join(parent_dir, os.path.basename(filepath) + '.json')
    return rrpath


# 获取转换后 OCR 结果的文件路径
def get_ocr_converted_result_path(filepath):
    # 获取用于存储转换后 OCR 结果的父目录路径
    parent_dir = fsutils.get_cache_path_for_category("ocr_converted_result")
    # 构建最终的转换后 OCR 结果文件路径
    crpath = os.path.join(parent_dir, os.path.basename(filepath) + '.json')
    return crpath


# 提取表格数据的函数
def extract_table_data(project_id: str, location: str, processor_id: str, filepath: str):
    # 获取原始 OCR 结果文件路径
    orrpath = get_ocr_raw_result_path(filepath)
    # 如果原始 OCR 结果文件已存在，则跳过扫描并打印相应信息
    if os.path.exists(orrpath):
        print("{} 跳过扫描，有之前生成的扫描文本".format(filepath))
        return

    # 创建 Document AI 客户端
    client = documentai.DocumentProcessorServiceClient()

    # 打开图像文件并读取其内容
    with open(filepath, 'rb') as image:
        image_content = image.read()

    # 设置图像的 MIME 类型
    mime_type = "image/jpeg"

    # 构建 Document AI 处理器的名称
    name = client.processor_path(project_id, location, processor_id)
    # 创建原始文档对象
    raw_document = documentai.RawDocument(
        content=image_content, mime_type=mime_type)

    # 构建请求对象
    request = {
        "name": name,
        "raw_document": raw_document,
    }

    # 向 Document AI 服务提交文档处理请求
    result = client.process_document(request=request)

    # 将结果转换为 JSON 字符串
    result_json_str = json.dumps(MessageToDict(result._pb), indent=2)

    # 将 JSON 字符串保存到文件
    with open(orrpath, 'w') as json_file:
        json_file.write(result_json_str)

    print("{} GCP扫描完成".format(orrpath))


# 后处理函数
def postprocess(filepath):
    # 获取原始 OCR 结果文件路径
    orrpath = get_ocr_raw_result_path(filepath)
    jsonstring = None
    with open(orrpath) as file:
        jsonstring = "".join(file.readlines())

    result_json = json.loads(jsonstring)

    # 提取文本和边界框，以 pytesseract.image_to_boxes 格式保存
    boxes = []

    # 读取图像文件并获取其高度、宽度和通道数
    image = cv2.imread(filepath)
    height, width, channel = image.shape

    for page in result_json.get('document', {}).get('pages', []):
        for line in page.get('tokens', []):
            # 获取文本行的顶点坐标
            vertices = line.get('layout').get('boundingPoly').get('normalizedVertices')
            points = []
            for i in vertices:
                points.append([int(i.get('x', 0) * width),
                               int(i.get('y', 0) * height)])
            points = np.array(points, dtype=np.int32)
            # 获取文本行的边界框坐标
            (x, y, w, h) = cv2.boundingRect(points)
            # 获取文本内容
            se = line.get('layout').get('textAnchor').get('textSegments')[0]
            start = int(se.get('startIndex', 0))
            end = int(se.get('endIndex'))
            chars = result_json.get('document').get('text')[start:end]
            if h / w > 5:
                # 如果文本行的高宽比大于 5，则将矩形在垂直方向上分成多个部分
                for j in range(len(chars)):
                    boxes.append(
                        [chars[j], x, y + h / len(chars) * j, w, h / len(chars)])
            else:
                # 否则直接添加整个文本行及其边界框
                boxes.append([chars, x, y, w, h])

    # 将提取的文本和边界框信息转换为 JSON 字符串
    boxes_str = json.dumps(boxes)

    # 获取转换后 OCR 结果文件路径
    ocrcpath = get_ocr_converted_result_path(filepath)
    with open(ocrcpath, 'w') as file:
        file.write(boxes_str)

    print("{} GCP转换完成".format(ocrcpath))


# 执行 OCR 过程的函数
def ocr(filepath):
    # 执行表格数据提取操作
    extract_table_data(project_id, location, processor_id, filepath)
    try:
        # 尝试执行后处理操作
        postprocess(filepath)
    except Exception as e:
        print("{} GCP OCR后处理失败".format(filepath))
        raise e


